package org.bouncycastle.cms.gm;

import org.bouncycastle.asn1.*;
import org.bouncycastle.asn1.cms.AttributeTable;
import org.bouncycastle.asn1.cms.ContentInfo;
import org.bouncycastle.asn1.cms.EncryptedContentInfo;
import org.bouncycastle.asn1.cms.EnvelopedData;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.cms.CMSEnvelopedData;
import org.bouncycastle.cms.CMSException;
import org.bouncycastle.cms.CMSTypedData;
import org.bouncycastle.cms.RecipientInfoGenerator;
import org.bouncycastle.operator.GenericKey;
import org.bouncycastle.operator.OutputEncryptor;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.util.HashMap;
import java.util.Iterator;

public class CMSEnvelopedDataGenerator extends CMSEnvelopedGenerator {
    
    public CMSEnvelopedDataGenerator() {
    }
    
    private CMSEnvelopedData doGenerate(CMSTypedData content, OutputEncryptor contentEncryptor) throws CMSException {
        ASN1EncodableVector recipientInfos = new ASN1EncodableVector();
        AlgorithmIdentifier encAlgId;
        ASN1OctetString     encContent;
        
        ByteArrayOutputStream bOut = new ByteArrayOutputStream();
        
        try {
            OutputStream cOut = contentEncryptor.getOutputStream(bOut);
            content.write(cOut);
            cOut.close();
        } catch (IOException e) {
            throw new CMSException("", e);
        }
        
        byte[] encryptedContent = bOut.toByteArray();
        
        encAlgId = contentEncryptor.getAlgorithmIdentifier();
        
        encContent = new BEROctetString(encryptedContent);
        
        GenericKey encKey = contentEncryptor.getKey();
        
        for (Iterator it = recipientInfoGenerators.iterator(); it.hasNext(); ) {
            RecipientInfoGenerator recipient = (RecipientInfoGenerator) it.next();
            recipientInfos.add(recipient.generate(encKey));
        }
        
        EncryptedContentInfo eci = new EncryptedContentInfo(content.getContentType(), encAlgId, encContent);
        
        ASN1Set unprotectedAttrSet = null;
        if (unprotectedAttributeGenerator != null) {
            AttributeTable attrTable = unprotectedAttributeGenerator.getAttributes(new HashMap());
            
            unprotectedAttrSet = new BERSet(attrTable.toASN1EncodableVector());
        }
        
        EnvelopedData envelopedData = new EnvelopedData(originatorInfo, new DERSet(recipientInfos), eci, unprotectedAttrSet);
        ContentInfo contentInfo = new ContentInfo(GMObjectIdentifiers.sm2_cms_envelopedData, envelopedData);
        return new CMSEnvelopedData(contentInfo);
    }
    
    /**
     * generate an enveloped object that contains an CMS Enveloped Data object using the given provider.
     *
     * @param content          the content to be encrypted
     * @param contentEncryptor the symmetric key based encryptor to encrypt the content with.
     */
    public CMSEnvelopedData generate(CMSTypedData content, OutputEncryptor contentEncryptor) throws CMSException {
        return doGenerate(content, contentEncryptor);
    }
    
}
